//
//  CustomURLProtocol.m
//  TEST_NSURLProtocol
//
//  Created by lianglibao on 2019/3/11.
//  Copyright © 2019年 lianglibao. All rights reserved.
//

#import "CustomURLProtocol.h"
#import "FetchIPAdressManager.h"
#import "NetWorkManager.h"

static NSString * const URLProtocolHandledKey = @"URLProtocolHandledKey";

@interface CustomURLProtocol ()<NSURLSessionTaskDelegate>
@property (nonatomic, strong) NSMutableDictionary <NSString*, NSURLSessionDataTask *> *dataTaskMap;
@property (nonatomic, strong) NSMutableArray <NSURLSessionDataTask *> *finishDataTaskArr;
@property (nonatomic, strong) dispatch_semaphore_t semaphore;
@property (nonatomic, strong) NSURLSessionTask *task;
@property (nonatomic, assign) BOOL isAddNotification;
@end

@implementation CustomURLProtocol

- (void)dealloc {
	[[NSNotificationCenter defaultCenter] removeObserver:self];
}

// MARK: - overwrite method
+ (BOOL)canInitWithRequest:(NSURLRequest *)request {
    if ([request.URL.scheme isEqualToString:@"http"] || [request.URL.scheme isEqualToString:@"https"]) {
        NSString *pathExtension = request.URL.pathExtension;
        // 不拦截(忽略)类似图片,字体资源请求.
        if ([ignorePathExtension() containsObject:[pathExtension lowercaseString]]) {
            if (request.URL.pathExtension.length) {
                NSLog(@"ignore url extension-->%@",request.URL.pathExtension);
            }
            return NO;
        }
        
        // 已经被替换为IP地址,则忽略..
        if (![FetchIPAdressManager isIPadress:request.URL.host] && request.URL.host.length) {
            return YES;
        } else {
            return NO;
        }
    }
    
    // 看看是否已经处理过了，防止无限循环
    if ([self.class propertyForKey:URLProtocolHandledKey inRequest:request]) {
        return NO;
    }
    
    return NO;
}

+ (NSURLRequest *)canonicalRequestForRequest:(NSURLRequest *)request {
    return request;
}

- (void)startLoading {
    // 解析DNS
    NSString *ip = [FetchIPAdressManager fetchIPAdressWithURL:self.request.URL];
    NSMutableURLRequest *request = [self.request mutableCopy];
    [request setValue:self.request.URL.host forHTTPHeaderField:@"host"];
	[request setURL:ip ? [NSURL URLWithString:ip] : self.request.URL];
	// 标记当前传入的Request已经被拦截处理过，
	// 防止在最开始又继续拦截处理
	[self.class setProperty:@(YES) forKey:URLProtocolHandledKey inRequest:request];

	// 处理请求.
    NetWorkManager *mgr = [NetWorkManager shareInstance];
	__weak typeof(self) weakSelf = self;
	NSURLSessionDataTask *dataTask = [mgr.manager dataTaskWithRequest:request
													   uploadProgress:nil
													 downloadProgress:nil
													completionHandler:nil];

	[mgr.manager setDataTaskDidReceiveResponseBlock:^NSURLSessionResponseDisposition(NSURLSession * _Nonnull session, NSURLSessionDataTask * _Nonnull dataTask, NSURLResponse * _Nonnull response) {
        [weakSelf.client URLProtocol:weakSelf didReceiveResponse:response cacheStoragePolicy:NSURLCacheStorageAllowed];
        return NSURLSessionResponseAllow;
    }];

	[mgr.manager setDataTaskDidReceiveDataBlock:^(NSURLSession * _Nonnull session, NSURLSessionDataTask * _Nonnull dataTask, NSData * _Nonnull data) {
        [weakSelf.client URLProtocol:weakSelf didLoadData:data];
    }];
	[dataTask resume];
		
	// 注册请求完成回调
	if (!self.isAddNotification) {
		dispatch_async(dispatch_get_main_queue(), ^{
			self.isAddNotification = YES;
			[[NSNotificationCenter defaultCenter] addObserver:self selector:@selector(taskFinished:) name:AFNetworkingTaskDidCompleteNotification object:nil];
		});
	}
	
	// 保存task
	NSString *taskIdentifier = [@(dataTask.taskIdentifier) stringValue];
	if (dataTask && taskIdentifier.length) {
		dispatch_semaphore_wait(self.semaphore, DISPATCH_TIME_FOREVER);
		[self.dataTaskMap setObject:dataTask forKey:taskIdentifier];
		dispatch_semaphore_signal(self.semaphore);
	}
}

- (void)stopLoading {
	dispatch_semaphore_wait(self.semaphore, DISPATCH_TIME_FOREVER);
	[self.finishDataTaskArr enumerateObjectsUsingBlock:^(NSURLSessionDataTask * _Nonnull obj, NSUInteger idx, BOOL * _Nonnull stop) {
		[obj cancel];
	}];
	[self.finishDataTaskArr removeAllObjects];
	dispatch_semaphore_signal(self.semaphore);
}

// MARK: - action
- (void)taskFinished:(NSNotification *)notice {
	 NSURLSessionDataTask *dataTask = notice.object;
	if (dataTask) {
		[self.finishDataTaskArr addObject:dataTask];
	}
	
	NSDictionary *userInfo = notice.userInfo;
	NSError *error = userInfo[AFNetworkingTaskDidCompleteErrorKey];
	if (error) {
		NSLog(@"%@", error);
		[self.client URLProtocol:self didFailWithError:error];
	} else {
		[self.client URLProtocolDidFinishLoading:self];
	}
}

// MARK: - getter
static inline NSSet *ignorePathExtension() {
    return [[NSSet alloc] initWithObjects:@"tif", @"tiff", @"jpg", @"jpeg", @"gif", @"png", @"ico", @"bmp", @"cur", @"apng", @"webp", @"woff", @"otf", nil];
}

- (NSMutableDictionary<NSString *,NSURLSessionDataTask *> *)dataTaskMap {
	if (!_dataTaskMap) {
		_dataTaskMap = [NSMutableDictionary dictionary];
	}
	return _dataTaskMap;;
}

- (NSMutableArray<NSURLSessionDataTask *> *)finishDataTaskArr {
	if (!_finishDataTaskArr) {
		_finishDataTaskArr = [NSMutableArray array];
	}
	return _finishDataTaskArr;
}

- (dispatch_semaphore_t)semaphore {
	if (!_semaphore) {
		_semaphore = dispatch_semaphore_create(1);
	}
	return _semaphore;
}

@end
